package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"net"
	"ntrace/enviroments"
	"ntrace/utils/routines"
	"time"
)

type ICMPSender struct {
	BaseSender
}

func NewICMPSender(ctx context.Context, task *TraceTask, stat *TraceStat) *ICMPSender {
	return &ICMPSender{
		BaseSender: BaseSender{
			Ctx:  ctx,
			Name: "ICMPSender",
			Task: task,
			Stat: stat,
		},
	}
}

func (s *ICMPSender) StartSender() error {
	routines.GoWithRecover("ICMPSender", s.doSend)
	return nil
}

func (s *ICMPSender) doSend() {
	rc, err := OpenSenderIP4RawPacketConn("ip4:icmp", "0.0.0.0")
	if err != nil {
		return
	}
	s.Conn = rc

	task := s.Task
	var seq uint16
	payload := AllocTxBuf(task.MsgLen)
	for turn := 0; turn < task.Count; turn++ {
		for _, srcIP := range task.SrcIPs {
			for ttl := uint8(1); ttl <= task.MaxTTL; ttl++ {
				select {
				case <-s.Ctx.Done():
					s.StopSender()
					return
				default:
					seq++
					_ = s.sendOnePkt(turn, srcIP, seq, ttl, payload)
				}
			}

			time.Sleep(time.Duration(task.FlowIntervalMs) * time.Millisecond)
		}
	}

	// 发送结束
	s.StopSender()
	TrigExit(s.Task.TimeoutMs)
}

func (s *ICMPSender) sendOnePkt(turn int, srcIP string, seq uint16, ttl uint8, payload []byte) error {
	task := s.Task
	target := task.DstIP
	pid := enviroments.PID
	golog.Debugf("send, id: %d, seq: %d", pid, ttl)
	txTs := time.Now().UnixNano()
	binary.BigEndian.PutUint16(payload[8:10], uint16(pid))
	binary.BigEndian.PutUint16(payload[10:12], seq)
	binary.BigEndian.PutUint64(payload[12:20], uint64(txTs))
	msg, err := EncodeIPICMPEchoRequestPacket(srcIP, target, task.TOS, ttl, uint16(pid), seq, payload)
	if err != nil {
		return err
	}
	err = s.Conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	if err != nil {
		return err
	}
	_, err = s.Conn.WriteToIP(msg, &net.IPAddr{IP: net.ParseIP(target)})
	if err != nil {
		return err
	}

	s.record(turn, srcIP, ttl, seq, txTs)
	return nil
}

func (s *ICMPSender) record(turn int, srcIP string, ttl uint8, seq uint16, txTs int64) {
	key := fmt.Sprintf("%s-%d", s.Task.DstIP, seq)
	pktRecord := &PktRecord{
		Key:   key,
		Turn:  turn,
		SrcIP: srcIP,
		DstIP: s.Task.DstIP,
		TTL:   ttl,

		TxTs: txTs,
	}

	s.Stat.PktRecordBucket.Add(pktRecord)
}
